import requests
import json
import time
from utils.image import download_image, mkdir, upload_image
import swapper
from init_scenes import scene_list
from loguru import logger
from comfy.comfyclient import ComfyClient

# SERVER_URL = 'http://127.0.0.1:7001/'
SERVER_URL = 'http://39.100.75.110/'

# 轮询间隔
SLEEP_TIME = 2

comfy_client = ComfyClient()

def start_connect():
  pay_task()
  
def post(url, params = {}, retry = 0):
  try:
    response = requests.post(SERVER_URL + url + '?ignore-csrf=aimagine-csrf', data=params)
    res_json = json.loads(response.text)
  except requests.RequestException as r:
    logger.error(r)
    post(url, params, retry + 1)
  except json.JSONDecodeError as e:
    return None
  
  return res_json

def update_task_progress(taskId, progress):
  """
  Purpose: 
  """
  post('common/service/update-task-progress', { "taskId": taskId, "progress": progress })
# end def

def pay_task():
  while True:
    try:
      res_json = post('common/service/pay-task')
      if (res_json == None):
        continue

      print(res_json)

      if res_json['success'] == True and res_json['task'] != None:
        task = res_json['task']
        swapper_task(task)
        
    except Exception as e:
      print(e)
    time.sleep(SLEEP_TIME)

def swapper_task(task):
  """
  换脸主函数
  """
  task_id = task['id']
  face = task['faceImage']
  scene_id = task['sceneId']
  face_file = download_image(face)
  logger.info('face download success ' + face)
  output_file = 'output_images/' + str(scene_id)  + '/' + task_id + '/' 
  target_dir = 'scene/' + str(scene_id) + '/'
  
  mkdir(output_file)
  output_results = []
  
  # 换脸流程
  scene_photos = scene_list[scene_id]
  for index, photo in enumerate(scene_photos):
    output_img = output_file + photo
    swapper.swap_face(face_file, target_dir + photo, output_img)
    logger.info('swap success ' + face + ' ' + photo)
    # 上传图片
    upload_image(output_img, output_img)
    print(task_id, index)
    # 修改进度
    update_task_progress(task_id, str(index + 1))
    logger.info('face upload success ' + output_img)
    output_results.append(output_img)

  # 修改状态
  post('common/service/update-task', {
    "taskId": task_id,
    "status": 2,
    "result": output_results
  })
  logger.info('update status success results--', output_results)

def faceIdWorkflow(json):
  """
  faceId文生图工作流
  用于漫画等场景
  """
  task = json['task']
  task_id = task['id']
  face = task['faceImage']
  prompt = task['prompt']
